import torch
import torch.nn as nn
import torch.nn.functional as F
from loss.hole_vaild import Hole, Valid
from loss.style import Style
from loss.total_variation import TotalVariation 
from loss.vgg import VGG

class Loss(nn.modules.loss._Loss):
    def __init__(self, opt):
        super(Loss, self).__init__()
        self.opt = opt
        self.loss = []
        if opt.hole > 0:
            self.loss.append({'loss': Hole(), 'weight': opt.hole})
        if opt.vaild > 0: 
            self.loss.append({'loss': Valid(), 'weight': opt.hole})
        if opt.perceptual > 0:
            self.loss.append({'loss': VGG(), 'weight': opt.perceptual})
        if opt.style > 0:
            self.loss.append({'loss': Style(), 'weight': opt.style})
        if opt.total_variation > 0:
            self.loss.append({'loss': TotalVariation(), 'weight': opt.total_variation})


    def forward(self, outs, comps, gts, mask):
        loss = 0 
        for l, w in self.loss:
            if l['loss_type'] == 'hole':
                loss += l['weight'] * l['loss_function'](outs, gts, mask)
            elif l['loss_type'] == 'vaild':
                loss += l['weight'] * l['loss_function'](outs, gts, 1-mask)
            elif l['loss_type'] == 'total_variation':
                loss += l['weight'] * l['loss_function'](comps)
            else:
                loss += l['weight'] * l['loss_function'](outs, gts) + l['weight'] * l['loss_function'](comps, gts)
            
        return loss
            
